import os, sys

parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 4)))
sys.path.insert(0, parent_path)

import traceback
import numpy as np
import cv2
from typing import *
from task.modules.processors.base_processor import BaseProcessor
from core.utils.visualize import get_color_map_list
from core.utils.PIL_draw import pic_text

class Visualization(BaseProcessor):
    def __init__(self, keys, vis_text_size=20) -> None:
        self.status = 0
        self.keys = keys
        self.text_size = vis_text_size
        self.color_list = get_color_map_list(80)
    
    def init_check(self):
        assert "input_data" in self.keys
        assert "in" in self.keys

    def __call__(self, data:Dict) -> Dict:
        batch_img = data[self.keys["input_data"]].copy()
        det_result = data[self.keys["in"]]
        batch_objects = []
        batch_vis = []
        # start_idx = 0
        for img_idx, bboxes in enumerate(det_result["boxes"]):
            draw_thickness = min(batch_img[img_idx].shape[:2]) // 320
            img_objects = []
            img = batch_img[img_idx].copy()
            for bbox in bboxes:
                cls_id, score, x1,y1,x2,y2 = bbox
                cls_id = int(cls_id)
                object = {
                    "box": [cls_id, score, x1, y1, x2, y2],
                    "cls_name": str(cls_id)
                }
                img_objects.append(object)
                # vis
                img = cv2.rectangle(img,(int(x1), int(y1)),(int(x2),int(y2)),color=self.color_list[cls_id],thickness=draw_thickness)
                text = str(cls_id)
                img = pic_text(img,text,(int(x1), int(y1)),self.color_list[cls_id],draw_thickness*10)
            batch_objects.append(img_objects)
            batch_vis.append(img)
        data[self.keys["out"]] = batch_vis[0]
        return data


if __name__ == "__main__":
    print("")
